import logging
import asyncio
import json

from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.requests import Request
from starlette.types import ASGIApp

from http_base import HttpResp

logger = logging.getLogger(__name__)


def init_cors_middleware(app: FastAPI):
    """初始化 CORS（跨域资源共享）中间件"""
    from env_config import settings
    app.add_middleware(
        CORSMiddleware,
        allow_origins=json.loads(settings.cors_allow_origins),
        allow_headers=['*'],
        allow_methods=['*'],
        max_age=3600
    )


class TimeoutMiddleware(BaseHTTPMiddleware):
    def __init__(self, app: ASGIApp, timeout: int = 15):
        super().__init__(app)
        self.timeout = timeout

    async def dispatch(self, request, call_next):
        try:
            return await asyncio.wait_for(call_next(request), timeout=self.timeout)
        except asyncio.TimeoutError:
            return JSONResponse(
                {'code': HttpResp.SYSTEM_TIMEOUT_ERROR.code, 'msg': HttpResp.SYSTEM_TIMEOUT_ERROR.msg, 'data': []})


def init_timeout_middleware(app: FastAPI):
    """初始化 超时处理 中间件"""
    from env_config import settings
    app.add_middleware(
        TimeoutMiddleware,
        timeout=settings.request_timeout)


def init_https_middleware(app: FastAPI):
    from dependencies.database import db

    @app.middleware("http")
    async def log_middleware(request: Request, call_next):
        logger.info('Request to access ' + request.url.path)
        try:
            return await call_next(request)
        finally:
            logger.info('Completed accessed ' + request.url.path)

    @app.middleware("http")
    async def db_connection_middleware(request: Request, call_next):
        # 检查数据库是否已连接
        if not db.is_connected:
            await db.connect()
        async with db.connection() as conn:
            request.state.conn = conn
            response = await call_next(request)
            return response

